package study.eggs.async.mdc.executor;

import org.slf4j.MDC;

import java.util.Map;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

public class MdcThreadPoolExecutor extends ThreadPoolExecutor {

//    public MdcThreadPoolExecutor(int corePoolSize,
//                                 int maximumPoolSize,
//                                 long keepAliveTime,
//                                 TimeUnit unit,
//                                 BlockingQueue<Runnable> workQueue) {
//        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue);
//    }
//
//    public MdcThreadPoolExecutor(int corePoolSize,
//                                 int maximumPoolSize,
//                                 long keepAliveTime,
//                                 TimeUnit unit,
//                                 BlockingQueue<Runnable> workQueue,
//                                 ThreadFactory threadFactory) {
//        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory);
//    }

//
    public MdcThreadPoolExecutor(int corePoolSize, int maximumPoolSize,
                                 long keepAliveTime, TimeUnit unit,
                                 BlockingQueue<Runnable> workQueue,
                                 ThreadFactory threadFactory,
                                 RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory, handler);
    }

    @Override
    public void execute(Runnable runnable) {
        //获取父线程 MDC 中的内容，必须在 run 方法之前，
        //否则等异步线程执行的时候有可能 MDC 里面的值已经被清空了，这个时候就会返回 null
        Map<String, String> context = MDC.getCopyOfContextMap();
        super.execute(() -> run(runnable, context));
    }

    private void run(Runnable runnable, Map<String, String> context) {
        //将父线程的 MDC 内容传给子线程
        MDC.setContextMap(context);
        try {
            //执行异步操作
            runnable.run();
        } finally {
            //清空 MDC 内容
            MDC.clear();
        }
    }
}
